import numpy as np
import random

# 神经元激活函数选择
class sigmoid:
    """
    sigmoid激活函数
    """
    def forward(self, z):
        """
        正向传播表达式
        z:加权输入θ@X
        """
        return 1.0 / (1.0 + np.exp(-z))

    def backward(self, output):
        """
        output:就是当前层的输出g(z)即是y
        :return:返回该激活函数的偏导数dg(z)/dz
        """
        return output * (1.0-output)


"""
损失函数选择
对于不同损失函数，仅仅代价函数J关于输出层加权输入z的偏导数dJ/dZ不同（输出层误差项不同），
而隐藏层的误差项均是由前一误差项递推而来，公式相同。
"""
class Square_J:
    # 平方损失函数
    def cost(self, label, output):
        # 单项损失函数
        return 0.5 * np.sum((label-output)**2)

    def d_outlayer(self, label, output):
        """
        求代价函数J关于输出层的输出[y=g(z)]的偏导数
        :return: 偏导数值
        """
        return -(label - output)


class CrossEntropy_J:
    # 交叉熵损失函数（为避免0带给log与分数到来的数值问题，我们加入0的近似值1e-5）
    def cost(self, label, output):
        # 单项损失函数
        return np.sum(-label*np.log(output+1e-5) - (1-label)*np.log(1-output+1e-5))

    def d_outlayer(self, label, output):
        """
           求代价函数J关于输出层的输出[y=g(z)]的偏导数
           :return: 偏导数值
        """
        return (-label/(output+1e-5)) + ((1-label)/(1-output+1e-5))


class FullConnectedLayer:
    """
    全连接层实现类：实例对象为神经网络的某一层，对该层进行正向传播、反向传播（计算误差项）、梯度下降等操作
    """
    def __init__(self, input_size, output_size, activator):
        """
        input_size: 当前层神经元输入列向量的维度（即上一层神经元个数）
        output_size: 当前层神经元的输出列向量的维度（即当前层神经元个数）
        activator: 选择不同激活函数
        """
        self.input_size = input_size
        self.output_size = output_size
        self.activator = activator
        self.W = np.random.uniform(-0.1, 0.1, (output_size, input_size))    # 初始化传递给当前层的权重参数
        self.b = np.zeros((self.output_size, 1))                            # 初始化传递给当前层的偏置项（当前层神经元个数的列向量）
        self.input = []                                                     # 当前层输入（上一层神经元输出）
        self.output = []                                                    # 当前层输出
        self.delta = []                                                     # 上一层层神经元的误差项δ（注意这里是上一层的误差项δ）
        self.gradient_W = []                                                # 传递给当前层的权重梯度矩阵（即dJ/dW）
        self.gradient_b = []                                                # 传递给当前层的偏置项梯度列向量（即dJ/db）

    def forward(self, input):
        """
        正向传播函数：更新属性：当前层的输入输出
        input: 当前层神经元的输入（上一层神经元的输出）
        """
        z = np.dot(self.W, input) + self.b                                  # 计算加权输入:z=W@X+b <=> z=θ@X
        self.input = input
        self.output = self.activator.forward(z)

    def backward(self, delta):
        """
        反向传播函数：计算上一层的误差项（因为我们计算隐藏层误差项是通过递推，而输出层误差项我们在模型里单独计算）
        delta: 后一层神经元的误差项向量
        """
        self.delta = self.activator.backward(self.input) * (self.W.T @ delta)   # 递推公式计算上一层的误差项δ
        self.gradient_W = delta @ self.input.T                                  # 更新传递给当前层权重梯度矩阵
        self.gradient_b = delta                                                 # 更新传递给当前层偏置项梯度矩阵

    def gradient_descent(self, rate):
        """
        梯度下降函数：实现简单的梯度下降算法（高级优化算法有兴趣可以自行实现）
        rate:学习率

        """
        self.W = self.W - rate*self.gradient_W                                  # 更新传递给当前层的权重矩阵
        self.b -= rate * self.gradient_b                                        # 更新传递给当前层的偏置项矩阵


class Neural_Network:
    """
    神经网络模型：实现训练、预测、计算损失、梯度检查等功能
    """
    def __init__(self, layers_struct, J):
        """
        layers_struct:神经网络层结构，每层神经元个数保存在数组中，例如[3, 5, 5, 4]
        J:选择代价函数：1.平方损失函数；2.交叉熵损失函数
        """
        self.rate = 0.033       # 学习率
        self.EPOCH = 10         # 训练轮数
        self.layers = []        # 保存该模型各全连接层对象
        self.J = J              # 模型选择的损失函数
        # 根据输入的神经网络结构初始化各全连接层
        for i in range(len(layers_struct)-1):
            self.layers.append(FullConnectedLayer(layers_struct[i], layers_struct[i+1], sigmoid()))

    def predict(self, sample):
        """
        预测指定单个样本
        sample: 单个样本
        :return: 预测结果（layer.output）
        """
        output = sample
        # 进行正向传播
        for layer in self.layers:
            layer.forward(output)
            output = layer.output
        return output

    def cal_gradient(self, label):
        """
        计算模型各层梯度
        label:当前样本的标签
        """
        # 单独计算输出层误差项δ，然后从后到前递推计算隐藏层误差项δ
        delta = (self.layers[-1].activator.backward(self.layers[-1].output) * self.J.d_outlayer(label, self.layers[-1].output))
        for layer in self.layers[::-1]:
            layer.backward(delta)
            delta = layer.delta

    def update_weight(self):
        # 更新模型各层系数矩阵（W、b）
        for layer in self.layers:
            layer.gradient_descent(self.rate)

    def cost(self, label, output):
        # 计算单个样本的损失值
        return self.J.cost(label, output)

    def train_one_sample(self, sample, label):
        """
        训练单个样本
        sample:单样本属性
        label:单样本标签
        """
        self.predict(sample)                            # 正向传播更新各层输入输出
        self.cal_gradient(label)                        # 反向传播更新各层参数的梯度
        self.update_weight()                            # 梯度下降更新各层参数

    def train(self, dataset, labels):
        """
        训练模型
        dataset:数据集特征（应当是三维数组，单个样本的特征用列向量表示）
        labels: 数据集标签（应当是三维数组，单个样本的标签用列向量0/1表示类别，序号与数据集特征一一对应）
        """
        for i in range(self.EPOCH):
            cost = 0
            print(f"正在进行第{i+1}轮训练:")
            for j in range(len(dataset)):
                self.train_one_sample(dataset[j], labels[j])
                cost += self.cost(labels[j], self.predict(dataset[j]))
            print(f"第{i+1}轮训练已完成，损失值J={cost/len(dataset)}")         # 计算模型损失值

    def check_gradient(self, sample, label):
        """
        梯度检查函数J-W，注意检查完毕后要关闭该函数。梯度检查独立于模型训练之外，可以单独运行
        sample:单个数据
        label:数据标签
        print:输出期望梯度与实际梯度的差值，观察梯度检查是否正确
        """
        # 首先正向传播一遍获得各层输入输出，再反向传播计算各层梯度
        self.predict(sample)
        self.cal_gradient(label)
        epsilon = 10e-4                                         # 设置极小项ε
        # 逐一计算每一层每一个参数的梯度是否正确
        for layer in self.layers:
            for i in range(layer.W.shape[0]):
                for j in range(layer.W.shape[1]):
                    layer.W[i, j] += epsilon
                    output = self.predict(sample)               # 更新当前层的输出
                    J2 = self.cost(label, output)               # 根据设置的w系数计算代价函数J(w+ε)
                    layer.W[i, j] -= 2 * epsilon
                    output = self.predict(sample)               # 更新当前层的输出
                    J1 = self.cost(label, output)               # 根据设置的w系数计算代价函数J(w-ε)
                    except_gradient = (J2-J1) / (2*epsilon)     # 计算期望梯度
                    print(f"对W进行梯度检查有:except_gradient - actual_gradient = {except_gradient-layer.gradient_W[i, j]}")

    def accuracy(self, test_X, test_y):
        """
        计算测试集准确率
        :param test_X: 测试数据集特征
        :param test_y: 测试数据集标签
        :return      : 返回模型准确率
        """
        n = 0
        for i in range(len(test_X)):
            predict = self.predict(test_X[i])
            #print(f"第{i+1}组预测:\n{test_y[i]}\n --> \n{predict}")
            if predict.argmax() == test_y[i].argmax():
                n += 1
        return n / len(test_X)

    def save_model(self, filename):
        """
        保存模型参数至txt文件中
        :param filename: 文件路径
        """
        with open(filename, 'a') as f:
            f.truncate(0)  # 每次重新保存模型时，删除历史数据
            for layer in self.layers:
                np.savetxt(f, np.c_[layer.W], fmt='%.6e', delimiter='\t')   # 保存权重矩阵用6位科学计数法表示，用\t分隔
                np.savetxt(f, np.c_[layer.b], fmt='%.6e', delimiter='\t')   # 保存偏置项列向量用6位科学计数法表示，用\t分隔

    def load_model(self, filename):
        """
        从txt文件中读取加载模型参数
        :param filename: 文件路径
        """
        with open(filename, 'r') as f:
            lines = f.readlines()                                           # 将txt文件中的所有数据按行以字符串形式存入lines数组
            for i in range(len(lines)):
                lines[i] = lines[i].split('\t')                             # 按\t分隔参数
                lines[i][-1] = lines[i][-1].replace('\n', '')               # 去掉每行最后一个元素的换行符\n
                for j in range(len(lines[i])):
                    lines[i][j] = np.float(lines[i][j])                     # 将所有字符串类型的元素转换为与模型参数相同类型的np.float
            last = 0
            for layer in self.layers:   # 加载各层神经元参数
                layer.W = np.array(lines[last:last+layer.output_size]).reshape((layer.output_size, layer.input_size))
                layer.b = np.array(lines[last+layer.output_size:last+layer.output_size*2]).reshape((layer.output_size, 1))
                last = last + layer.output_size*2

